def translateToAcid(seq):
    def StriOrderlyByNum(stri, num):
        if len(stri) % num == 0:
            return stri
        else:
            redundantNum = len(stri) % num
            totalNum = len(stri)
            return stri[0:totalNum - redundantNum]

    table = {
        'ATA': 'I', 'ATC': 'I', 'ATT': 'I', 'ATG': 'M',
        'ACA': 'T', 'ACC': 'T', 'ACG': 'T', 'ACT': 'T',
        'AAC': 'N', 'AAT': 'N', 'AAA': 'K', 'AAG': 'K',
        'AGC': 'S', 'AGT': 'S', 'AGA': 'R', 'AGG': 'R',
        'CTA': 'L', 'CTC': 'L', 'CTG': 'L', 'CTT': 'L',
        'CCA': 'P', 'CCC': 'P', 'CCG': 'P', 'CCT': 'P',
        'CAC': 'H', 'CAT': 'H', 'CAA': 'Q', 'CAG': 'Q',
        'CGA': 'R', 'CGC': 'R', 'CGG': 'R', 'CGT': 'R',
        'GTA': 'V', 'GTC': 'V', 'GTG': 'V', 'GTT': 'V',
        'GCA': 'A', 'GCC': 'A', 'GCG': 'A', 'GCT': 'A',
        'GAC': 'D', 'GAT': 'D', 'GAA': 'E', 'GAG': 'E',
        'GGA': 'G', 'GGC': 'G', 'GGG': 'G', 'GGT': 'G',
        'TCA': 'S', 'TCC': 'S', 'TCG': 'S', 'TCT': 'S',
        'TTC': 'F', 'TTT': 'F', 'TTA': 'L', 'TTG': 'L',
        'TAC': 'Y', 'TAT': 'Y', 'TAA': '_', 'TAG': '_',
        'TGC': 'C', 'TGT': 'C', 'TGA': '_', 'TGG': 'W',
    }

    seq1 = StriOrderlyByNum(seq, 3)
    protein = ""
    if len(seq1) % 3 == 0:
        for i in range(0, len(seq1), 3):
            codon = seq1[i:i + 3]
            ToAcid = table[codon]
            if ToAcid != "_":
                protein += ToAcid
            else:
                break
    else:
        protein = "-"
        print "please check seq num!", len(seq1) % 3

    return protein


def AcidDeqExtNucSeq(AcidMatSeq,AcidTotSeq,NucTotSeq):
    Pos=AcidTotSeq.find(AcidMatSeq)
    NucMatSeq = "-"
    if Pos > -1:
        NulTotLen=len(NucTotSeq)
        AcidMatLen=len(AcidMatSeq)
        if (Pos + AcidMatLen)*3 <= NulTotLen:
            NucMatSeq = NucTotSeq[Pos*3:(Pos+AcidMatLen)*3]
    return NucMatSeq

def TCRTRVseqExt(seq, kseq, CutDisMatNum):
    def getStrAllFindPos(source, key, posfrac):
        sourceLen = len(source)
        NeedBigPos = int(sourceLen * posfrac)

        PosList = []
        findPos = source.find(key)
        while findPos >= 0:
            if findPos <= NeedBigPos:
                PosList.append(findPos)
            findPos = source.find(key, findPos + 1)
        if len(PosList) > 0:
            return PosList
        else:
            return False

    def seqmat2(seq1, seq2, CutNum):
        # input:
        # CutNum:two seq cut to the same to compare
        # output:
        # DisMatchNum(two seq dismatch num or false)
        # false(one more seq's length less than CutNum)

        L1 = len(seq1)
        L2 = len(seq2)

        if L1 < CutNum or L2 < CutNum:
            return False

        seq1 = seq1[0:CutNum]
        seq2 = seq2[0:CutNum]

        DisMatchNum = 0
        for i in range(0, CutNum):
            if seq1[i] != seq2[i]:
                DisMatchNum += 1
        return DisMatchNum

    def getSeqMisMatchNum(seq, kseq, PosList):
        # input:
        # seq:long seq
        # kseq:known seq
        # PosList:known seq in long seq possible Pos
        # output:
        # PosMisNumDict({Pos:MisMatchNum})(0,1,2,... or "-")

        PosMisNumDict = {}
        for i in PosList:
            seqExt = seq[i:]
            MisMat = seqmat2(seqExt, kseq, len(kseq))
            if MisMat >= -1:
                PosMisNumDict[i] = MisMat
            else:
                PosMisNumDict[i] = "-"
        return PosMisNumDict

    def ChoMatPos(DisMatNumDict, DisMatCuf):
        # input:
        # DisMatNumDict({Pos:DisMatchNum})
        # DisMatCuf(a value of max dismatch)
        # output:
        # Pos or False
        AllMatList = []
        for i in sorted(DisMatNumDict.items(), key=lambda x: x[1], reverse=False):
            Pos = i[0]
            DisMatNum = i[1]
            if DisMatNum == 0:
                AllMatList.append(Pos)
        # choose all match pos
        if len(AllMatList) > 0:
            AllMatList.sort()
            return AllMatList[0]
        # dismatch num is more low and pos is Front
        else:
            for i in sorted(DisMatNumDict.items(), key=lambda x: x[1], reverse=False):
                Pos = i[0]
                DisMatNum = i[1]

                if DisMatNum <= DisMatCuf:
                    return Pos
            return False

    # TCRTRVseqExt main
    # ATG
    HeadKSeq = kseq[0:3]
    PosList = getStrAllFindPos(seq, HeadKSeq, 0.75)

    if PosList != False:
        DisMatNumDict = getSeqMisMatchNum(seq, kseq, PosList)
        FinalPos = ChoMatPos(DisMatNumDict, CutDisMatNum)
        return seq[FinalPos:]
    else:
        return "-"


def getTRABVFa(inputFa, BFlag):
    import re
    geneFaDict = {}

    LineNum = -1
    GeneSeq = ""
    geneName = "NA"
    geneFaDict[geneName] = ""

    f = open(inputFa)
    line = f.readline()
    while line:
        line = line.strip("\n")
        m = re.match("(>)(\S*)", line)
        n = re.match("(\d+)(\s+)(\S+)", line)
        if m:
            geneFaDict[geneName] = GeneSeq  # save before Gene seq
            geneName = m.group(2)
            if geneName not in geneFaDict:  # a new Gene
                geneFaDict[geneName] = ""
                LineNum = 0
                GeneSeq = ""
        elif n:
            LineNum += 1
            line = line.replace(" ", "")
            line = line.replace("\t", "")
            s = re.match("(\d+)([A-Z]+)", line)
            if s:  # like:61ATCGATCG
                seq = s.group(2)
                if LineNum == 2:  # gene seq ATG start
                    if seq[0:3] == "ATG":
                        GeneSeq += seq
                    else:
                        print geneName, "warning:(this start not ATG)", line
                        GeneSeq += seq
                elif LineNum > 2:  # other line seq
                    GeneSeq += seq
            else:
                print "warning line(no num):", line
        else:
            print "this line can't see!"
        line = f.readline()
    f.close()

    if BFlag == 1:
        geneFaDict1 = {}
        for g in geneFaDict:
            gene = g.split("-")[0]
            if gene not in geneFaDict1:
                geneFaDict1[gene] = geneFaDict[g]
        return geneFaDict1
    return geneFaDict


if __name__ == '__main__':
    print ""
    # CDR3 = "CSALDRGQYTGELFF"
    # SeqFaAcid = "MLSLLLLLLGLGSVFSAVISQKPSRDICQRGTSLTIQCQVDSQVTMMFWYRQQPGQSLTLIATANQGSEATYESGFVIDKFPISRPNLTFSTLTVSNMSPEDSSIYLCSALDRGQYTGELFFGEGSRLTVLEDLKNVFPPEVAVFEPS"
    # SeqFaExt = "ATGCTGAGTCTTCTGCTCCTTCTCCTGGGACTAGGCTCTGTGTTCAGTGCTGTCATCTCTCAAAAGCCAAGCAGGGATATCTGTCAACGTGGAACCTCCCTGACGATCCAGTGTCAAGTCGATAGCCAAGTCACCATGATGTTCTGGTACCGTCAGCAACCTGGACAGAGCCTGACACTGATCGCAACTGCAAATCAGGGCTCTGAGGCCACATATGAGAGTGGATTTGTCATTGACAAGTTTCCCATCAGCCGCCCAAACCTAACATTCTCAACTCTGACTGTGAGCAACATGAGCCCTGAAGACAGCAGCATATATCTCTGCAGCGCTCTGGACAGGGGGCAGTACACCGGGGAGCTGTTTTTTGGAGAAGGCTCTAGGCTGACCGTACTGGAGGACCTGAAAAACGTGTTCCCACCCGAGGTCGCTGTGTTTGAGCCATCAGA"
    # CDR3Dna = AcidDeqExtNucSeq(CDR3, SeqFaAcid, SeqFaExt)
    # print CDR3Dna
    """
    TCR-1249	TCR-1249	P_Tumor_TTCGTGCTCATCTACTTAAC	TRBV6-1:None:TRBJ2-7:CASSEMVGSSYEQYF	GTTCCCCTATCACCGATGCACAGACCCAGAAGACCCCTCCATCCTGTAGCACCTGCCATGAGCATCGGGCTCCTGTGCTGTGTGGCCTTTTCTCTCCTGTGGGCAAGTCCAGTGAATGCTGGTGTCACTCAGACCCCAAAATTCCAGGTCCTGAAGACAGGACAGAGCATGACACTGCAGTGTGCCCAGGATATGAACCATAACTCCATGTACTGGTATCGACAAGACCCAGGCATGGGACTGAGGCTGATTTATTACTCAGCTTCTGAGGGTACCACTGACAAAGGAGAAGTCCCCAATGGCTACAATGTCTCCAGATTAAACAAACGGGAGTTCTCGCTCAGGCTGGAGTCGGCTGCTCCCTCCCAGACATCTGTGTACTTCTGTGCCAGCAGTGAAATGGTGGGAAGCTCCTACGAGCAGTACTTCGGGCCGGGCACCAGGCTCACGGTCACAGAGGGCCTGAAAAACGTGT	TRBV6-1	None	TRBJ2-7	CASSEMVGSSYEQYF	B	ATGAGCATCGGGCTCCTGTGCTGTGTGGCCTTTTCTCTCCTGTGGGCAAGTCCAGTGAATGCTGGTGTCACTCAGACCCCAAAATTCCAGGTCCTGAAGACAGGACAGAGCATGACACTGCAGTGTGCCCAGGATATGAACCATAACTCCATGTACTGGTATCGACAAGACCCAGGCATGGGACTGAGGCTGATTTATTACTCAGCTTCTGAGGGTACCACTGACAAAGGAGAAGTCCCCAATGGCTACAATGTCTCCAGATTAAACAAACGGGAGTTCTCGCTCAGGCTGGAGTCGGCTGCTCCCTCCCAGACATCTGTGTACTTCTGTGCCAGCAGTGAAATGGTGGGAAGCTCCTACGAGCAGTACTTCGGGCCGGGCACCAGGCTCACGGTCACAGAGGGCCTGAAAAACGTGT	MSIGLLCCVAFSLLWASPVNAGVTQTPKFQVLKTGQSMTLQCAQDMNHNSMYWYRQDPGMGLRLIYYSASEGTTDKGEVPNGYNVSRLNKREFSLRLESAAPSQTSVYFCASSEMVGSSYEQYFGPGTRLTVTEGLKNV	-	-	TGTGCCAGCAGTGAAATGGTGGGAAGCTCCTACGAGCAGTACTTC
    """
    seq = "GTTCCCCTATCACCGATGCACAGACCCAGAAGACCCCTCCATCCTGTAGCACCTGCCATGAGCATCGGGCTCCTGTGCTGTGTGGCCTTTTCTCTCCTGTGGGCAAGTCCAGTGAATGCTGGTGTCACTCAGACCCCAAAATTCCAGGTCCTGAAGACAGGACAGAGCATGACACTGCAGTGTGCCCAGGATATGAACCATAACTCCATGTACTGGTATCGACAAGACCCAGGCATGGGACTGAGGCTGATTTATTACTCAGCTTCTGAGGGTACCACTGACAAAGGAGAAGTCCCCAATGGCTACAATGTCTCCAGATTAAACAAACGGGAGTTCTCGCTCAGGCTGGAGTCGGCTGCTCCCTCCCAGACATCTGTGTACTTCTGTGCCAGCAGTGAAATGGTGGGAAGCTCCTACGAGCAGTACTTCGGGCCGGGCACCAGGCTCACGGTCACAGAGGGCCTGAAAAACGTGT"
    VGene = "TRBV6-1"
    VGene2 = VGene.split("-")[0]

    TRABVFaFile = "/home/hsmy/code/ik/TCR/singlecell/TCRPriDesign/DB/VGeneTRAV.fa"
    VGeneFaDict = getTRABVFa(TRABVFaFile, 1)
    VGeneFaDict2 = getTRABVFa(TRABVFaFile, 0)

    if VGene in VGeneFaDict2:
        StandSeq = VGeneFaDict2[VGene]
    else:
        StandSeq = VGeneFaDict[VGene2]

    extSeq = TCRTRVseqExt(seq, StandSeq[0:20], 5)
    acid = translateToAcid(extSeq)
    print acid
